import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import argparse
import os
from datetime import datetime

# 设置中文字体
try:
    # 尝试使用系统已安装的中文字体
    plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
    plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
except Exception as e:
    print(f"设置中文字体时出错: {e}")
    print("将继续使用默认字体")

def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(
        description='可视化细胞系-化合物AUC矩阵数据，按疾病类型着色'
    )
    parser.add_argument('--input', '-i', 
                       default='d:/playground/cellauc/cell_line_compound_auc_matrix.csv',
                       help='输入CSV文件路径，默认为 cell_line_compound_auc_matrix.csv')
    parser.add_argument('--output', '-o',
                       help='输出图形路径（可选）')
    parser.add_argument('--palette', '-p', 
                       default='husl',
                       help='颜色方案，支持：husl, tab10, Set2, viridis, plasma 或 custom，默认为husl')
    parser.add_argument('--figsize', '-f', 
                       default='12,8',
                       help='图形大小，格式为"宽,高"，默认为"12,8"')
    parser.add_argument('--plot-type', '-t',
                       default='box',
                       choices=['scatter', 'box', 'bar'],
                       help='绘图类型：散点图(scatter)、箱线图(box)或条形图(bar)，默认为带有散点数据的箱线图')
    parser.add_argument('--cell-line-col',
                       default='cell_line',
                       help='细胞系列名，默认为 cell_line')
    parser.add_argument('--disease-col',
                       default='disease',
                       help='疾病列名，默认为 disease')
    
    return parser.parse_args()

def load_and_preprocess_data(input_file, cell_line_col, disease_col):
    """
    加载并预处理数据
    将矩阵格式转换为长格式，提取化合物列和疾病信息
    """
    # 读取数据
    df = pd.read_csv(input_file)
    print(f"原始数据形状: {df.shape}")
    print(f"数据列名: {list(df.columns)}")
    
    # 验证必需的列是否存在
    if cell_line_col not in df.columns:
        raise ValueError(f"错误：数据中找不到细胞系列 '{cell_line_col}'")
    if disease_col not in df.columns:
        raise ValueError(f"错误：数据中找不到疾病列 '{disease_col}'")
    
    # 识别基础信息列和化合物列
    info_columns = [cell_line_col, disease_col]
    # 添加 lineage 列（如果存在）
    if 'lineage' in df.columns:
        info_columns.append('lineage')
    
    # 提取化合物列 - 假设不是基础信息列的就是化合物列
    compound_columns = [col for col in df.columns if col not in info_columns]
    print(f"识别到 {len(compound_columns)} 个化合物列: {compound_columns[:10]}")  # 只显示前10个
    if len(compound_columns) > 10:
        print(f"... 等 {len(compound_columns) - 10} 个化合物")
    
    # 将数据转换为长格式
    long_df = pd.melt(
        df, 
        id_vars=info_columns, 
        value_vars=compound_columns, 
        var_name='compound', 
        value_name='auc_value'
    )
    
    # 移除NaN值（如果有）
    initial_rows = len(long_df)
    long_df = long_df.dropna(subset=['auc_value'])
    removed_rows = initial_rows - len(long_df)
    if removed_rows > 0:
        print(f"移除了 {removed_rows} 行含NaN的AUC值数据")
    
    print(f"转换后的数据形状: {long_df.shape}")
    
    # 获取唯一疾病类型
    disease_types = long_df[disease_col].unique()
    print(f"疾病类型 ({len(disease_types)}): {disease_types}")
    
    # 重命名疾病列以保持一致性
    if disease_col != 'disease':
        long_df = long_df.rename(columns={disease_col: 'disease'})
    
    return long_df, compound_columns

def get_palette_colors(palette_name, group_count):
    """
    获取与cellauc.py兼容的颜色方案
    
    参数:
    palette_name (str): 调色板名称，支持 'husl', 'tab10', 'Set2', 'viridis', 'plasma' 或 'custom'
    group_count (int): 分组数量
    
    返回:
    list: 颜色列表
    """
    # 自定义颜色列表 - 与cellauc.py保持一致
    custom_colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', 
                     '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf']
    
    if palette_name == 'custom':
        # 如果自定义颜色不足，则循环使用
        if group_count <= len(custom_colors):
            colors = custom_colors[:group_count]
        else:
            colors = custom_colors * (group_count // len(custom_colors) + 1)
            colors = colors[:group_count]
    else:
        # 使用seaborn内置调色板
        try:
            colors = sns.color_palette(palette_name, group_count)
        except ValueError:
            print(f"警告：未知的调色板 '{palette_name}'，使用默认的 'husl'")
            colors = sns.color_palette('husl', group_count)
    
    return colors

def create_visualization(long_df, compound_columns, palette_name, figsize, plot_type, output_path=None):
    """
    创建可视化图表
    """
    # 解析图形大小
    try:
        width, height = map(float, figsize.split(','))
    except ValueError:
        print(f"警告：无效的图形大小参数 '{figsize}'，使用默认值 12,8")
        width, height = 12, 8
    
    # 获取疾病类型数量并生成颜色映射
    disease_types = long_df['disease'].unique()
    colors = get_palette_colors(palette_name, len(disease_types))
    color_map = {disease: color for disease, color in zip(disease_types, colors)}
    
    # 创建图形
    fig, ax = plt.subplots(figsize=(width, height))
    
    # 根据绘图类型创建不同的图表
    if plot_type == 'scatter':
        # 散点图 - 每个点代表一个细胞系-化合物组合
        sns.scatterplot(
            data=long_df,
            x='compound',
            y='auc_value',
            hue='disease',
            palette=color_map,
            alpha=0.7,
            s=80,
            ax=ax
        )
    elif plot_type == 'box':
        # 箱线图 - 展示每个化合物在不同疾病类型中的分布
        sns.boxplot(
            data=long_df,
            x='compound',
            y='auc_value',
            hue='disease',
            palette=color_map,
            ax=ax
        )
        # 在箱线图上添加散点数据
        sns.stripplot(
            data=long_df,
            x='compound',
            y='auc_value',
            hue='disease',
            palette=color_map,
            dodge=True,
            size=4,
            edgecolor='auto',
            linewidth=0.5,
            ax=ax,
            legend=False
        )
        # 去除重复的图例项
        handles, labels = ax.get_legend_handles_labels()
        by_label = dict(zip(labels, handles))
        ax.legend(by_label.values(), by_label.keys(), title='疾病类型', bbox_to_anchor=(1.05, 1), loc='upper left')
    elif plot_type == 'bar':
        # 条形图 - 展示每个化合物在不同疾病类型中的平均值
        sns.barplot(
            data=long_df,
            x='compound',
            y='auc_value',
            hue='disease',
            palette=color_map,
            errorbar='sd',
            ax=ax
        )
        # 在条形图上添加散点数据
        sns.stripplot(
            data=long_df,
            x='compound',
            y='auc_value',
            hue='disease',
            palette=color_map,
            dodge=True,
            size=4,
            edgecolor='auto',
            linewidth=0.5,
            ax=ax,
            legend=False
        )
        # 去除重复的图例项
        handles, labels = ax.get_legend_handles_labels()
        by_label = dict(zip(labels, handles))
        ax.legend(by_label.values(), by_label.keys(), title='疾病类型', bbox_to_anchor=(1.05, 1), loc='upper left')
    
    # 设置标题和标签
    ax.set_title('化合物AUC值按疾病类型着色', fontsize=16, fontweight='bold')
    ax.set_xlabel('化合物', fontsize=12)
    ax.set_ylabel('AUC值', fontsize=12)
    
    # 调整x轴标签角度以避免重叠
    plt.xticks(rotation=45, ha='right')
    
    # 调整图例位置
    plt.legend(title='疾病类型', bbox_to_anchor=(1.05, 1), loc='upper left')
    
    # 添加网格线以便于阅读
    ax.grid(True, linestyle='--', alpha=0.7)
    
    # 紧凑布局
    plt.tight_layout()
    
    # 如果指定了输出路径，保存图形
    if output_path:
        os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
        plt.savefig(output_path, dpi=300, bbox_inches='tight')
        print(f"图形已保存至: {output_path}")
    else:
        # 显示图形
        plt.show()

def main():
    """主函数"""
    try:
        # 解析命令行参数
        args = parse_arguments()
        
        # 验证输入文件是否存在
        if not os.path.exists(args.input):
            print(f"错误：输入文件不存在: {args.input}")
            return
        
        # 加载并预处理数据
        print("正在加载和预处理数据...")
        long_df, compound_columns = load_and_preprocess_data(
            args.input,
            args.cell_line_col,
            args.disease_col
        )
        
        # 创建可视化
        print(f"正在创建可视化... 使用调色板: {args.palette}")
        create_visualization(
            long_df,
            compound_columns,
            args.palette,
            args.figsize,
            args.plot_type,
            args.output
        )
        
        print("可视化完成！")
    except Exception as e:
        print(f"发生错误: {str(e)}")
        import traceback
        print("详细错误信息:")
        traceback.print_exc()

if __name__ == "__main__":
    main()